#!/usr/bin/env python3

import argparse
import json
import os
import signal
import socket
import ssl
import sys
import threading
import time
from typing import Dict, List, Optional

from colorama import Fore, Style, init as colorama_init
from websocket import WebSocketApp


def parse_headers(header_list: List[str]) -> Dict[str, str]:
    headers: Dict[str, str] = {}
    for item in header_list or []:
        if ":" not in item:
            raise ValueError(f"无效的请求头格式: {item}. 正确格式示例: 'Authorization: Bearer xxx'")
        name, value = item.split(":", 1)
        headers[name.strip()] = value.strip()
    return headers


def human_time() -> str:
    return time.strftime("%H:%M:%S", time.localtime())


class WebSocketTester:
    def __init__(
        self,
        url: str,
        headers: Optional[Dict[str, str]] = None,
        subprotocols: Optional[List[str]] = None,
        insecure: bool = False,
        origin: Optional[str] = None,
        timeout: Optional[float] = None,
        ping_interval: Optional[float] = 25.0,
        ping_timeout: Optional[float] = 20.0,
        show_frames: bool = True,
    ) -> None:
        self.url = url
        # build headers list; include Origin if provided
        merged_headers: Dict[str, str] = dict(headers or {})
        if origin:
            merged_headers.setdefault("Origin", origin)
        self.headers = [f"{k}: {v}" for k, v in merged_headers.items()]
        self.subprotocols = subprotocols
        self.insecure = insecure
        self.origin = origin
        self.timeout = timeout
        self.ping_interval = ping_interval
        self.ping_timeout = ping_timeout
        self.show_frames = show_frames
        self._app: Optional[WebSocketApp] = None
        self._rx_thread: Optional[threading.Thread] = None
        self._closed = threading.Event()
        self._errored: Optional[str] = None

    def _sslopt(self) -> Dict[str, object]:
        if self.url.startswith("wss://"):
            if self.insecure:
                return {
                    "cert_reqs": ssl.CERT_NONE,
                    "check_hostname": False,
                }
            return {
                "cert_reqs": ssl.CERT_REQUIRED,
            }
        # 对于 ws:// 协议，明确设置为 None 避免库的自动检测
        return None

    def _on_open(self, _ws):
        print(f"{Fore.GREEN}[{human_time()}] 已连接 -> {self.url}{Style.RESET_ALL}")

    def _on_message(self, _ws, message):
        if isinstance(message, bytes):
            prefix = f"[{human_time()}] 收到二进制({len(message)} bytes):"
            print(f"{Fore.CYAN}{prefix}{Style.RESET_ALL}")
            if self.show_frames:
                print(message)
        else:
            prefix = f"[{human_time()}] 收到文本:"
            print(f"{Fore.CYAN}{prefix}{Style.RESET_ALL} {message}")

    def _on_error(self, _ws, error):
        self._errored = str(error)
        print(f"{Fore.RED}[{human_time()}] 错误: {error}{Style.RESET_ALL}")

    def _on_close(self, _ws, status_code, msg):
        print(
            f"{Fore.YELLOW}[{human_time()}] 连接已关闭 code={status_code} reason={msg}{Style.RESET_ALL}"
        )
        self._closed.set()

    def connect(self) -> None:
        # 尝试解析域名为 IP 地址，避免 websocket-client 的域名处理问题
        try:
            from urllib.parse import urlparse
            parsed = urlparse(self.url)
            if parsed.hostname and not parsed.hostname.replace('.', '').isdigit():
                # 是域名，尝试解析为 IP
                ip = socket.gethostbyname(parsed.hostname)
                resolved_url = self.url.replace(parsed.hostname, ip)
                actual_url = resolved_url
            else:
                actual_url = self.url
        except Exception as e:
            # 域名解析失败，使用原始 URL
            actual_url = self.url
        
        self._app = WebSocketApp(
            actual_url,
            header=self.headers,
            on_open=self._on_open,
            on_message=self._on_message,
            on_error=self._on_error,
            on_close=self._on_close,
            subprotocols=self.subprotocols,
        )

        def run():
            self._app.run_forever(
                sslopt=self._sslopt(),
                ping_interval=self.ping_interval,
                ping_timeout=self.ping_timeout,
            )

        self._rx_thread = threading.Thread(target=run, daemon=True)
        self._rx_thread.start()

        # wait for open or error or timeout
        start = time.time()
        while not self._closed.is_set():
            if self._errored:
                raise RuntimeError(self._errored)
            if self.timeout is not None and (time.time() - start) > self.timeout:
                raise TimeoutError("连接超时")
            if self._app and self._app.sock and self._app.sock.connected:
                return
            time.sleep(0.05)

    def send_text(self, text: str) -> None:
        if not self._app or not self._app.sock or not self._app.sock.connected:
            raise RuntimeError("尚未连接")
        self._app.send(text)
        print(f"{Fore.MAGENTA}[{human_time()}] 已发送文本:{Style.RESET_ALL} {text}")

    def send_bytes(self, data: bytes) -> None:
        if not self._app or not self._app.sock or not self._app.sock.connected:
            raise RuntimeError("尚未连接")
        self._app.send(data, opcode=0x2)
        print(f"{Fore.MAGENTA}[{human_time()}] 已发送二进制({len(data)} bytes){Style.RESET_ALL}")

    def close(self) -> None:
        if self._app:
            try:
                self._app.close()
            finally:
                self._closed.wait(timeout=2.0)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="WebSocket 测试小工具 (基于 websocket-client)",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--url",
        default="wss://echo.websocket.events",
        help="WebSocket 地址，支持 ws:// 与 wss://",
    )
    print(f'parser.add_argument:')
    parser.add_argument(
        "-H",
        "--header",
        action="append",
        default=[],
        help="附加请求头，格式: 'Name: Value'，可多次指定",
    )
    parser.add_argument(
        "--subprotocol",
        action="append",
        default=None,
        help="Sec-WebSocket-Protocol 子协议，可多次指定",
    )
    parser.add_argument(
        "--origin",
        default=None,
        help="Origin 头",
    )
    parser.add_argument(
        "--insecure",
        action="store_true",
        help="跳过 TLS 证书校验(仅 wss://)",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=10.0,
        help="连接超时秒数",
    )
    parser.add_argument(
        "--ping-interval",
        type=float,
        default=25.0,
        help="保持心跳的 ping 间隔秒数",
    )
    parser.add_argument(
        "--ping-timeout",
        type=float,
        default=20.0,
        help="心跳 ping 超时秒数",
    )
    parser.add_argument(
        "-m",
        "--message",
        default=None,
        help="一次性发送文本消息，若不指定则进入交互模式",
    )
    parser.add_argument(
        "--binary",
        action="store_true",
        help="将 --message 按十六进制字符串发送为二进制，如 'ff00aa'",
    )
    parser.add_argument(
        "-n",
        "--count",
        type=int,
        default=1,
        help="发送次数(配合 --message 使用)",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="以 JSON 格式美化打印文本消息",
    )
    return parser


def main(argv: Optional[List[str]] = None) -> int:
    colorama_init(autoreset=True)
    args = build_parser().parse_args(argv)
    print(f'args.url: {args.url}')

    try:
        headers = parse_headers(args.header)
    except ValueError as e:
        print(str(e))
        return 2

    tester = WebSocketTester(
        url=args.url,
        headers=headers,
        subprotocols=args.subprotocol,
        insecure=bool(args.insecure),
        origin=args.origin,
        timeout=args.timeout,
        ping_interval=args.ping_interval,
        ping_timeout=args.ping_timeout,
    )

    def handle_sigint(_sig, _frm):
        print(f"{Fore.YELLOW}\n[{human_time()}] 收到中断信号，正在关闭...{Style.RESET_ALL}")
        tester.close()
        sys.exit(130)

    signal.signal(signal.SIGINT, handle_sigint)

    try:
        tester.connect()
    except Exception as e:
        print(f"{Fore.RED}连接失败: {e}{Style.RESET_ALL}")
        return 1

    # one-shot send
    if args.message is not None:
        for _ in range(max(1, args.count)):
            if args.binary:
                try:
                    data = bytes.fromhex(args.message.replace(" ", ""))
                except ValueError:
                    print(f"{Fore.RED}--binary 模式需要十六进制字符串，如 'ff00aa'{Style.RESET_ALL}")
                    tester.close()
                    return 2
                tester.send_bytes(data)
            else:
                msg = args.message
                if args.json:
                    try:
                        msg_obj = json.loads(msg)
                        msg = json.dumps(msg_obj, ensure_ascii=False)
                    except json.JSONDecodeError:
                        pass
                tester.send_text(msg)

        # 给对端一些响应时间
        time.sleep(1.0)
        tester.close()
        return 0

    # interactive mode
    print(
        f"{Fore.WHITE}进入交互模式: 输入文本并回车发送。输入 /quit 退出，/hex ff00aa 发送二进制。{Style.RESET_ALL}"
    )
    while True:
        try:
            line = input("")
        except EOFError:
            break
        except KeyboardInterrupt:
            break

        if not line:
            continue
        if line.strip().lower() in {"/quit", ":q", "exit"}:
            break
        if line.strip().startswith("/hex "):
            hexstr = line.strip()[5:].replace(" ", "")
            try:
                data = bytes.fromhex(hexstr)
            except ValueError:
                print(f"{Fore.RED}无效的十六进制字符串{Style.RESET_ALL}")
                continue
            tester.send_bytes(data)
            continue
        if args.json:
            try:
                obj = json.loads(line)
                line = json.dumps(obj, ensure_ascii=False)
            except json.JSONDecodeError:
                pass
        tester.send_text(line)

    tester.close()
    return 0


if __name__ == "__main__":
    sys.exit(main())
